#include <cusp/precond/spai.h>
#include <cusp/krylov/bicgstab.h>
#include <cusp/gallery/poisson.h>
#include <cusp/gallery/eye.h>
#include <cusp/csr_matrix.h>

#include <cusp/frobenius_norm.h>
#include <cusp/print.h>
#include <cusp/multiply.h>
#include <cusp/elementwise.h>

#include <iostream>

template <typename Monitor>
void report_status(Monitor& monitor)
{
    if (monitor.converged())
    {
        std::cout << "Solver converged to " << monitor.tolerance() << " tolerance";
        std::cout << " after " << monitor.iteration_count() << " iterations";
        std::cout << " (" << monitor.residual_norm() << " final residual)" << std::endl;
    }
    else
    {
        std::cout << "Solver reached iteration limit " << monitor.iteration_limit() << " before converging";
        std::cout << " to " << monitor.tolerance() << " tolerance ";
        std::cout << " (" << monitor.residual_norm() << " final residual)" << std::endl;
    }
}

int main(int argc, char** argv)
{
    typedef int                 IndexType;
    typedef float               ValueType;
    typedef cusp::device_memory MemorySpace;

    // create an empty sparse matrix structure
    cusp::csr_matrix<IndexType, ValueType, MemorySpace> A;

    if (argc < 2)
    {
        cusp::gallery::poisson5pt(A, 512, 512);
    }
    else
    {
        // read matrix from file
        cusp::io::read_matrix_market_file(A, std::string(argv[1]));
    }

    std::cout << "size " << A.num_rows << " x " << A.num_cols << std::endl;

    // solve without preconditioning
    {
        std::cout << "\nSolving with no preconditioner" << std::endl;
    
        // allocate storage for solution (x) and right hand side (b)
        cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 0);
        cusp::array1d<ValueType, MemorySpace> b(A.num_rows, 1);

        // set stopping criteria (iteration_limit = 1000, relative_tolerance = 1e-6)
        cusp::default_monitor<ValueType> monitor(b, 1000, 1e-6);
        
        // solve
        cusp::krylov::bicgstab(A, x, b, monitor);

        // report status
        report_status(monitor);
    }

    // solve SPAI preconditioner
    {
        std::cout << "\nSolving with SPAI preconditioner" << std::endl;
        
        // allocate storage for solution (x) and right hand side (b)
        cusp::array1d<ValueType, MemorySpace> x(A.num_rows, 0);
        cusp::array1d<ValueType, MemorySpace> b(A.num_rows, 1);

        // set stopping criteria (iteration_limit = 1000, relative_tolerance = 1e-6)
        cusp::default_monitor<ValueType> monitor(b, 1000, 1e-6);

        // setup preconditioner
	cusp::csr_matrix<IndexType, ValueType, MemorySpace> I;
	cusp::gallery::eye(I, A.num_rows, A.num_cols);
        cusp::precond::spai<ValueType, MemorySpace> M(A, A);

        // solve
        cusp::krylov::bicgstab(A, x, b, monitor, M);
        
        // report status
        report_status(monitor);

        // report norms
        cusp::csr_matrix<int,float,cusp::device_memory> P;
        cusp::multiply(A, M.matrix, P);
        cusp::subtract(P, I, P);
        std::cout << "Fnorm(A)\t" << cusp::frobenius_norm(A) << "\n";
        std::cout << "Fnorm(AM-I)\t" << cusp::frobenius_norm(P) << std::endl;
    }

    return 0;
}

